#include "mt3d_cuda.h"

void mt3d::solve()
{
    // Create handles
	cublasHandle_t cubHandle;
	cusparseHandle_t cusHandle;

	cublasCreate(&cubHandle);
	cusparseCreate(&cusHandle);

	// Allocate GPU memory & copy matrix/vector to device
	cudaMalloc(&d_A_, nz_num_ * sizeof(cuDoubleComplex));
	cudaMalloc(&d_rowIdxA_, nz_num_ * sizeof(int));
	cudaMalloc(&d_rowPtrA_, (edge_num_ + 1) * sizeof(int));
	cudaMalloc(&d_colIdxA_, nz_num_ * sizeof(int));

	cudaMalloc(&d_P_, edge_num_ * sizeof(cuDoubleComplex));
	cusparseCreateDnVec(&dvec_P_, edge_num_, d_P_, CUDA_C_64F);

	cudaMemcpy(d_A_, h_A_.get(), nz_num_ * sizeof(cuDoubleComplex), cudaMemcpyHostToDevice);
	cudaMemcpy(d_rowIdxA_, h_rowIdxA_.get(), nz_num_ * sizeof(int), cudaMemcpyHostToDevice);
	cudaMemcpy(d_colIdxA_, h_colIdxA_.get(), nz_num_ * sizeof(int), cudaMemcpyHostToDevice);

    // Convert matrix A from COO format to CSR format
	cusparseXcoo2csr(cusHandle, d_rowIdxA_, nz_num_, edge_num_, d_rowPtrA_, CUSPARSE_INDEX_BASE_ZERO);

	// Create sparse matrix
	cusparseCreateCsr(&smat_A_, edge_num_, edge_num_, nz_num_, d_rowPtrA_, d_colIdxA_, d_A_, CUSPARSE_INDEX_32I,
		CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, CUDA_C_64F);

	// This is just used to get bufferSize;
	cuda_cd *d_tmp;
	cudaMalloc(&d_tmp, edge_num_ * sizeof(cuDoubleComplex));

	cusparseDnVecDescr_t dvec_tmp;
	cusparseCreateDnVec(&dvec_tmp, edge_num_, d_tmp, CUDA_C_64F);

	cusparseSpMV_bufferSize(cusHandle, CUSPARSE_OPERATION_NON_TRANSPOSE, &_one, smat_A_,
		dvec_P_, &_zero, dvec_tmp, CUDA_C_64F, CUSPARSE_SPMV_ALG_DEFAULT, &bufferSize_);
	cudaMalloc(&d_buf_, bufferSize_);

	// Preconditioning using diagonal elements of the kernel matrix. Get the diagonal elemenets
	clcg_smZcsr_get_diagonal(d_rowPtrA_, d_colIdxA_, d_A_, edge_num_, d_P_);

	h_X_TE_.resize(edge_num_, _zero);
    h_X_TM_.resize(edge_num_, _zero);

    clcg_para my_para = clcg_default_parameters();
    my_para.epsilon = 1e-7;
    set_clcg_parameter(my_para);
    set_report_interval(1000);

	std::clog << "Solving the PDE in TE mode ... \n";
	//Minimize(cubHandle, cusHandle, h_X_TE_.get(), h_B_TE_.get(), edge_num_, nz_num_, CLCG_BICG_SYM);
	MinimizePreconditioned(cubHandle, cusHandle, h_X_TE_.get(), h_B_TE_.get(), edge_num_, nz_num_, CLCG_PCG);

	std::clog << "Solving the PDE in TM mode ... \n";
	//Minimize(cubHandle, cusHandle, h_X_TM_.get(), h_B_TM_.get(), edge_num_, nz_num_, CLCG_BICG_SYM);
	MinimizePreconditioned(cubHandle, cusHandle, h_X_TM_.get(), h_B_TM_.get(), edge_num_, nz_num_, CLCG_PCG);

	// Free Device memory
	cudaFree(d_buf_);
	cudaFree(d_A_);
	cudaFree(d_rowIdxA_);
	cudaFree(d_rowPtrA_);
	cudaFree(d_colIdxA_);
	cudaFree(d_P_);
	cudaFree(d_tmp);

	cusparseDestroyDnVec(dvec_P_);
	cusparseDestroyDnVec(dvec_tmp);
	cusparseDestroySpMat(smat_A_);

	cublasDestroy(cubHandle);
	cusparseDestroy(cusHandle);
    return;
}
